# Code for the Molct model.
# Reference:
# Zhang, J.; Zhou, Y.; Lei, Y.-K.; Yang, Y. I.; Gao, Y. Q.,
# Molecular CT: unifying geometry and representation learning for molecules at different scales [J/OL].
# arXiv preprint, 2020: arXiv:2012.11816 [2020-12-22]. https://arxiv.org/abs/2012.11816

import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk

from typing import Optional, Union, Callable
from jax import nn as nn
from .interaction.molct_interaction import InteractionUnit
from ..utils.filter import get_filter
from ..utils.activation import get_activation

class MolCT(hk.Module):
    r"""Molecular CT model.

    ### Args:

        dim_feature (int): Dimension of node and edge vectors.

        dim_node_emb (int): Dimension of node embeddings. Default: None.

        dim_edge_emb (int): Dimension of edge embeddings. Default: None.

        is_edge_update (bool): Whether to update edge vectors. Default: False.

        is_coupled_interaction (bool): Whether to use coupled interaction. Default: False.

        n_interaction (int): Number of interaction units. Default: 3.

        n_heads (int): Number of heads in multi-head attention. Default: 8.

        dim_outer_pdct (int): Dimension of outer product. Default: 32.

        num_transition (int): Number of transition in hyperformer. Default: 4.

        fp_type (dtype): Floating point type. Default: jnp.float32.

        activation (Callable, str): Activation function, Default: 'silu'.

        name (str): Name of the module. Default: 'molct'.
    
    """

    def __init__(self,
                 dim_feature: int,
                 dim_node_emb: Optional[int] = None,
                 dim_edge_emb: Optional[int] = None,
                 is_edge_update: bool = False,
                 is_coupled_interaction: bool = False,
                 n_interaction: int = 3,
                 n_heads: int = 8,
                 dim_outer_pdct: int = 32,
                 num_transition: int = 4,
                 fp_type = jnp.float32,
                 activation: Optional[Union[str, Callable]] = 'silu',
                 name: str = "molct",
                 **kwargs,
                 ):
        super().__init__(name=name)

        self.dim_feature = dim_feature
        self.n_heads = n_heads
        self.dim_outer_pdct = dim_outer_pdct
        self.num_transition = num_transition
        self.fp_type = fp_type
        self.num_interaction = n_interaction
        self.is_edge_update = is_edge_update
        self.is_coupled_intaction = is_coupled_interaction
        self.activation = get_activation(activation)

        # build filter
        if dim_node_emb is None:
            dim_node_emb = dim_feature
        if dim_edge_emb is None:
            dim_edge_emb = dim_feature

        self.node_filter = None
        if dim_feature != dim_node_emb:
            self.node_filter = get_filter(cls_name='residual',
                                          dim_in=dim_node_emb,
                                          dim_out=dim_feature,
                                          activation=self.activation,
                                          name="node_filter")
        
        self.edge_filter = None
        if dim_feature != dim_edge_emb:
            self.edge_filter = get_filter(cls_name='residual',
                                          dim_in=dim_edge_emb,
                                          dim_out=dim_feature,
                                          activation=self.activation,
                                          name="edge_filter")
        
        # build interaction
        self.build_interaction()

    def build_interaction(self):

        if self.is_coupled_intaction:
            self.interaction = \
                [
                    InteractionUnit(dim_feature=self.dim_feature,
                                    n_heads=self.n_heads,
                                    dim_outer_pdct=self.dim_outer_pdct,
                                    num_transition=self.num_transition,
                                    is_edge_update=self.is_edge_update,
                                    activation=self.activation,
                                    fp_type=self.fp_type,
                                    name=f"interaction_unit") # type: ignore
                ] * self.num_interaction
            
        else:
            self.interaction = \
                [
                    InteractionUnit(dim_feature=self.dim_feature,
                                    n_heads=self.n_heads,
                                    dim_outer_pdct=self.dim_outer_pdct,
                                    num_transition=self.num_transition,
                                    is_edge_update=self.is_edge_update,
                                    activation=self.activation,
                                    fp_type=self.fp_type,
                                    name=f"interaction_unit_{_idx}")  # type: ignore
                    for _idx in range(self.num_interaction)
                ]

    
    def __call__(self,
                 node_emb: jnp.ndarray,
                 node_mask: jnp.ndarray,
                 edge_emb: jnp.ndarray,
                 edge_mask: jnp.ndarray,
                 edge_cutoff: jnp.ndarray,
                 ):
        r"""
        ### Args:

            node_emb (jnp.ndarray): Node embeddings. Shape: (A, F).

            node_mask (jnp.ndarray): Node masks. Shape: (A,).

            edge_emb (jnp.ndarray): Edge embeddings. Shape: (A, A, F).

            edge_mask (jnp.ndarray): Edge masks. Shape: (A, A).

            edge_cutoff (jnp.ndarray): Edge cutoffs. Shape: (A, A).
        
        """
        
        # (A, A)
        diagonal_mask = jnp.eye(node_emb.shape[-2], dtype=jnp.bool_)
        # (A, A) or (A, A) -> (A, A)
        edge_mask = jnp.logical_or(edge_mask, diagonal_mask)

        if self.node_filter is not None:
            node_vec = self.node_filter(node_emb)
        else:
            node_vec = node_emb

        if self.edge_filter is not None:
            edge_vec = self.edge_filter(edge_emb)
        else:
            edge_vec = edge_emb
        
        for interaction_net in self.interaction:
            node_vec, edge_vec = interaction_net(
                node_vec=node_vec,
                node_mask=node_mask,
                edge_vec=edge_vec,
                edge_mask=edge_mask,
                edge_cutoff=edge_cutoff,
                )
        
        return node_vec, edge_vec
    